#!/usr/bin/python3 -u
# Upload the container to a registry.  Note this
# is distinct from `upload-oscontainer` which
# only applies to (hopefully soon only older)
# versions of RHCOS but not FCOS.

import argparse
import json
import os
import tempfile
import shutil
import subprocess
import sys

cosa_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, cosa_dir)

from cosalib import cmdlib

transports = {
    'podman':'containers-storage:',
    'dir':'dir:',
    'registry':'docker://',
    'docker-archive':'docker-archive:',
    'docker':'docker-daemon:',
}

parser = argparse.ArgumentParser()
parser.add_argument("--authfile", help="Authentication file",
                    action='store')
parser.add_argument("--format", help="Image format for destination", choices=['oci', 'v2s2'], action='store')
parser.add_argument("--tag-suffix", metavar='SUFFIX', help="Append SUFFIX to container tag")
parser.add_argument("name", metavar='NAME[:TAG]', help="destination image reference")
parser.add_argument("--image", default='ostree', help="Container image to push", choices=['ostree', 'extensions-container'])
# add for nestos
parser.add_argument("-d","--dest-repository", help="The destination repository to push")
parser.add_argument("--no-tls-verify", help="Don't verify the destination repository TLS when specified",action="store_true")
parser.add_argument("--transport", help="destination image transport, support: \
                    podman, dir, registry, docker-archive, docker",
                    default='registry')

args = parser.parse_args()

with open('builds/builds.json') as f:
    builds = json.load(f)['builds']
if len(builds) == 0:
    cmdlib.fatal("No builds found")
latest_build = builds[0]['id']
arch = cmdlib.get_basearch()
latest_build_path = f"builds/{latest_build}/{arch}"

metapath = f"{latest_build_path}/meta.json"
with open(metapath) as f:
    meta = json.load(f)
ociarchive = os.path.join(latest_build_path, meta['images'][args.image]['path'])

skopeoargs = ['skopeo', 'copy']
if args.authfile is None:
    args.authfile = os.environ.get("REGISTRY_AUTH_FILE")
if args.authfile is not None:
    skopeoargs.extend(['--authfile', args.authfile])
if args.format is not None:
    skopeoargs.extend(['--format', args.format])
if ":" not in args.name:
    container_name = args.name
    # If a specific tag wasn't requested, add a default one
    container_tag = f"{latest_build}-{arch}"
else:
    # Strip the tag out, as we will be injecting the digest below
    # Note this implicitly errors out if there's more than one ':'
    container_name, container_tag = args.name.rsplit(':')
if args.tag_suffix:
    container_tag = f"{container_tag}-{args.tag_suffix}"
if args.no_tls_verify:
    skopeoargs.extend(["--insecure-policy"])
    skopeoargs.extend(["--dest-tls-verify=false"])
if args.transport not in transports.keys():
    print(f"Not supported transport:{args.transport}")
    exit(1)
with tempfile.NamedTemporaryFile(dir='tmp', prefix='push-container-digestfile') as df:
    skopeoargs.append(f"--digestfile={df.name}")
    # TODO: to test for private registry
    if args.dest_repository:
        skopeoargs.extend([f"oci-archive:{ociarchive}", f"{transports[args.transport]}{args.dest_repository}/{container_name}:{container_tag}"])
    else:
        skopeoargs.extend([f"oci-archive:{ociarchive}", f"{transports[args.transport]}{container_name}:{container_tag}"])
    print(subprocess.list2cmdline(skopeoargs))
    subprocess.check_call(skopeoargs)
    df.seek(0)
    digest = df.read().decode('utf-8').strip()
    # Inject the oscontainer with SHA256 into the build metadata
    container = 'base-oscontainer'
    if args.image != 'ostree':
        container = args.image
    meta[container] = {
        'image': container_name,
        'digest': digest
    }
    metapath_new = f"{metapath}.new"
    with open(metapath_new, 'w') as f:
        json.dump(meta, f, sort_keys=True)
    shutil.move(metapath_new, metapath)
